// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

#include "osstd.hxx"
#include <wincrypt.h>
#include <ntsecapi.h>

//
//  Computes the bit reversed CRC32C (Castagnoli) compatible with SSE4.2 CRC instruction
//
// The faster table version.  Use the init routine below with the reversed CRC polynomial 0x1EDC6F41 to create it
//

static const ULONG dwCRC32_LOOKUP_TABLE [256] =
{
     0x00000000, 0xF26B8303, 0xE13B70F7, 0x1350F3F4, 0xC79A971F, 0x35F1141C, 0x26A1E7E8, 0xD4CA64EB,
     0x8AD958CF, 0x78B2DBCC, 0x6BE22838, 0x9989AB3B, 0x4D43CFD0, 0xBF284CD3, 0xAC78BF27, 0x5E133C24,
     0x105EC76F, 0xE235446C, 0xF165B798, 0x030E349B, 0xD7C45070, 0x25AFD373, 0x36FF2087, 0xC494A384,
     0x9A879FA0, 0x68EC1CA3, 0x7BBCEF57, 0x89D76C54, 0x5D1D08BF, 0xAF768BBC, 0xBC267848, 0x4E4DFB4B,
     0x20BD8EDE, 0xD2D60DDD, 0xC186FE29, 0x33ED7D2A, 0xE72719C1, 0x154C9AC2, 0x061C6936, 0xF477EA35,
     0xAA64D611, 0x580F5512, 0x4B5FA6E6, 0xB93425E5, 0x6DFE410E, 0x9F95C20D, 0x8CC531F9, 0x7EAEB2FA,
     0x30E349B1, 0xC288CAB2, 0xD1D83946, 0x23B3BA45, 0xF779DEAE, 0x05125DAD, 0x1642AE59, 0xE4292D5A,
     0xBA3A117E, 0x4851927D, 0x5B016189, 0xA96AE28A, 0x7DA08661, 0x8FCB0562, 0x9C9BF696, 0x6EF07595,
     0x417B1DBC, 0xB3109EBF, 0xA0406D4B, 0x522BEE48, 0x86E18AA3, 0x748A09A0, 0x67DAFA54, 0x95B17957,
     0xCBA24573, 0x39C9C670, 0x2A993584, 0xD8F2B687, 0x0C38D26C, 0xFE53516F, 0xED03A29B, 0x1F682198,
     0x5125DAD3, 0xA34E59D0, 0xB01EAA24, 0x42752927, 0x96BF4DCC, 0x64D4CECF, 0x77843D3B, 0x85EFBE38,
     0xDBFC821C, 0x2997011F, 0x3AC7F2EB, 0xC8AC71E8, 0x1C661503, 0xEE0D9600, 0xFD5D65F4, 0x0F36E6F7,
     0x61C69362, 0x93AD1061, 0x80FDE395, 0x72966096, 0xA65C047D, 0x5437877E, 0x4767748A, 0xB50CF789,
     0xEB1FCBAD, 0x197448AE, 0x0A24BB5A, 0xF84F3859, 0x2C855CB2, 0xDEEEDFB1, 0xCDBE2C45, 0x3FD5AF46,
     0x7198540D, 0x83F3D70E, 0x90A324FA, 0x62C8A7F9, 0xB602C312, 0x44694011, 0x5739B3E5, 0xA55230E6,
     0xFB410CC2, 0x092A8FC1, 0x1A7A7C35, 0xE811FF36, 0x3CDB9BDD, 0xCEB018DE, 0xDDE0EB2A, 0x2F8B6829,
     0x82F63B78, 0x709DB87B, 0x63CD4B8F, 0x91A6C88C, 0x456CAC67, 0xB7072F64, 0xA457DC90, 0x563C5F93,
     0x082F63B7, 0xFA44E0B4, 0xE9141340, 0x1B7F9043, 0xCFB5F4A8, 0x3DDE77AB, 0x2E8E845F, 0xDCE5075C,
     0x92A8FC17, 0x60C37F14, 0x73938CE0, 0x81F80FE3, 0x55326B08, 0xA759E80B, 0xB4091BFF, 0x466298FC,
     0x1871A4D8, 0xEA1A27DB, 0xF94AD42F, 0x0B21572C, 0xDFEB33C7, 0x2D80B0C4, 0x3ED04330, 0xCCBBC033,
     0xA24BB5A6, 0x502036A5, 0x4370C551, 0xB11B4652, 0x65D122B9, 0x97BAA1BA, 0x84EA524E, 0x7681D14D,
     0x2892ED69, 0xDAF96E6A, 0xC9A99D9E, 0x3BC21E9D, 0xEF087A76, 0x1D63F975, 0x0E330A81, 0xFC588982,
     0xB21572C9, 0x407EF1CA, 0x532E023E, 0xA145813D, 0x758FE5D6, 0x87E466D5, 0x94B49521, 0x66DF1622,
     0x38CC2A06, 0xCAA7A905, 0xD9F75AF1, 0x2B9CD9F2, 0xFF56BD19, 0x0D3D3E1A, 0x1E6DCDEE, 0xEC064EED,
     0xC38D26C4, 0x31E6A5C7, 0x22B65633, 0xD0DDD530, 0x0417B1DB, 0xF67C32D8, 0xE52CC12C, 0x1747422F,
     0x49547E0B, 0xBB3FFD08, 0xA86F0EFC, 0x5A048DFF, 0x8ECEE914, 0x7CA56A17, 0x6FF599E3, 0x9D9E1AE0,
     0xD3D3E1AB, 0x21B862A8, 0x32E8915C, 0xC083125F, 0x144976B4, 0xE622F5B7, 0xF5720643, 0x07198540,
     0x590AB964, 0xAB613A67, 0xB831C993, 0x4A5A4A90, 0x9E902E7B, 0x6CFBAD78, 0x7FAB5E8C, 0x8DC0DD8F,
     0xE330A81A, 0x115B2B19, 0x020BD8ED, 0xF0605BEE, 0x24AA3F05, 0xD6C1BC06, 0xC5914FF2, 0x37FACCF1,
     0x69E9F0D5, 0x9B8273D6, 0x88D28022, 0x7AB90321, 0xAE7367CA, 0x5C18E4C9, 0x4F48173D, 0xBD23943E,
     0xF36E6F75, 0x0105EC76, 0x12551F82, 0xE03E9C81, 0x34F4F86A, 0xC69F7B69, 0xD5CF889D, 0x27A40B9E,
     0x79B737BA, 0x8BDCB4B9, 0x988C474D, 0x6AE7C44E, 0xBE2DA0A5, 0x4C4623A6, 0x5F16D052, 0xAD7D5351
};

#if 0
    // this is actually calculating the reverse CRC
    // computing the reverse CRC of 0 gives the table entry above
    //
    static uint32 calc_entry( uint32 crc, unsigned index, uint32 poly )
{
        for(INT i=0;i<8;i++)
    {
            if((crc ^ index) & 1)
                crc = (crc>>1)^poly;
            else
                crc>>= 1;
            index >>= 1;
    }
        return crc;
}

    static void initcrc32table( uint32 poly, uint32 CrcXorTable[256] )
{
        for (unsigned i = 0; i < 256; ++i) {
            uint32 entry = calc_entry( 0, i, poly );
            CrcXorTable[i] = entry;
    }
}
#endif

BOOL g_fProcessorSupportsCRC32 = fFalse;

#if (defined (_M_AMD64) || defined (_M_IX86))
#if defined (_M_AMD64)
#define COMPUTE_CRC_SIZE 8
#else
#define COMPUTE_CRC_SIZE 4
#endif

ULONG
ProcessorComputeCRC32(
    _In_reads_bytes_(Length) const  BYTE * Buffer,
                                    ULONG  Length
)
{
    ULONG  i;
    ULONG  crc;
    ULONG  endBytes;


    crc = 0xffffffff;

    // First do the bytes in groups of 4 or 8
    endBytes = Length % COMPUTE_CRC_SIZE;
    for ( i=0; i<Length/COMPUTE_CRC_SIZE; i++ )
    {
#if defined (_M_AMD64)
        crc = (ULONG) _mm_crc32_u64((unsigned __int64)crc, *(const unsigned __int64 UNALIGNED *)Buffer);
#else
        crc = _mm_crc32_u32(crc, *(const UINT UNALIGNED *)Buffer);
#endif
        Buffer += COMPUTE_CRC_SIZE;
    }

    // Calculate for the remaining bytes.
    for(i=0; i<endBytes; i++)
    {
        crc = _mm_crc32_u8(crc, *Buffer++);
    }
    crc = crc ^ 0xffffffff;

    return crc;
}
#endif // #if defined (_M_AMD64) || defined (_M_IX86)

void
OSInitializeProcessorSupportsCRC32()
{
    g_fProcessorSupportsCRC32 = fFalse;

#if (defined (_M_AMD64) || defined (_M_IX86))
    INT CPUInfo[4] = { 0, 0, 0, 0 };

    //
    // Get processor information.
    // use Infotype parameter as 1 to get the SSE 4 (CRC support_ information.
    // Refer msdn description of __cpuid for details.
    //
    __cpuid( CPUInfo, 1 );

    //
    // Bit 20 of 2nd array index indicates support for SSE 4.2 CRC32
    // instruction on the current platform.
    //
    if ( CPUInfo[2] & 0x100000 )
    {
        g_fProcessorSupportsCRC32 = fTrue;
    }
#endif
}

ULONG
Crc32Checksum(
    _In_reads_bytes_(cbData)    const   BYTE *pbData,
    _In_                                ULONG cbData )
{
    ULONG crcProcessor = 0;

#if (defined (_M_AMD64) || defined (_M_IX86))
    if ( g_fProcessorSupportsCRC32 )
    {
        //
        // Processor supports instruction to compute CRC32. 
        //
        crcProcessor = ProcessorComputeCRC32( pbData, cbData );
#ifndef DEBUG
        return crcProcessor;
#endif
        // On debug builds, compute both CRCs and compare them
    }
#endif

    ULONG crc = 0xffffffff;
    for (; cbData > 0; cbData-- )
    {
        crc = dwCRC32_LOOKUP_TABLE[ BYTE(crc ^ *pbData++) ] ^ (crc >> 8);
    }
    crc = (crc ^ 0xffffffff);

    if ( g_fProcessorSupportsCRC32 )
    {
        Assert( crc == crcProcessor );
        return crcProcessor;
    }
    else
    {
        return crc;
    }
}


HCRYPTPROV g_hAESProv = NULL;
CCriticalSection g_critAESProv( CLockBasicInfo( CSyncBasicInfo( _T( "g_critAESProv" ) ), rankAESProv, 0 ) );
#define BlockSizeAes256 16

BOOL FOSEncryptionPreinit()
{
    OSInitializeProcessorSupportsCRC32();
    return fTrue;
}

void OSEncryptionPostterm() {}

ERR
ErrOSEncryptionInit()
{
    return JET_errSuccess;
}

ERR
ErrOSIAESProviderInit()
{
    if ( g_hAESProv != NULL )
    {
        return JET_errSuccess;
    }

    g_critAESProv.Enter();

    if ( g_hAESProv != NULL )
    {
        g_critAESProv.Leave();
        return JET_errSuccess;
    }

    HCRYPTPROV hAESProv = NULL;
    if ( !CryptAcquireContextW( &hAESProv, NULL, MS_ENH_RSA_AES_PROV_W, PROV_RSA_AES, CRYPT_VERIFYCONTEXT ) )
    {
        g_critAESProv.Leave();
        return ErrOSErrFromWin32Err(GetLastError());
    }

    const VOID *hPrev = AtomicExchangePointer( (VOID **)&g_hAESProv, (VOID *)hAESProv );
    Assert( hPrev == NULL );
    g_critAESProv.Leave();
    return JET_errSuccess;
}

void
OSEncryptionTerm()
{
    if ( g_hAESProv != NULL )
    {
        CryptReleaseContext( g_hAESProv, 0 );
        g_hAESProv = NULL;
    }
}

#include <pshpack1.h>
// PERSISTED
struct AES256KEY
{
    BYTE                            Version;
    UnalignedLittleEndian<ULONG>    Checksum;
    BYTE                            pbKey[0];
};

ERR ErrOSEncryptionVerifyKey(
        _In_reads_bytes_(cbKey)                         const   BYTE *pbKey,
        _In_                                                    ULONG cbKey )
{
    AES256KEY *pKey = (AES256KEY *)pbKey;

    if ( cbKey < sizeof(AES256KEY) ||
         pKey->Version != JET_EncryptionAlgorithmAes256 ||
         pKey->Checksum != Crc32Checksum( pKey->pbKey, cbKey - sizeof(AES256KEY) ) )
    {
        return ErrERRCheck( JET_errInvalidParameter );
    }

    return JET_errSuccess;
}

ERR
ErrOSCreateAes256Key(
    _Out_writes_bytes_to_opt_(*pcbKeySize, *pcbKeySize) BYTE *pbKey,
    _Inout_                                             ULONG *pcbKeySize )
{
    ERR err = JET_errSuccess;
    HCRYPTKEY hKey = NULL;
    AES256KEY *pKey = (AES256KEY *)pbKey;
    ULONG cbKeySize = *pcbKeySize;

    CallR( ErrOSIAESProviderInit() );

    if ( !CryptGenKey( g_hAESProv, CALG_AES_256, CRYPT_EXPORTABLE, &hKey ) )
    {
        return ErrOSErrFromWin32Err(GetLastError());
    }

    // Prepend the version
    if ( cbKeySize >= sizeof(AES256KEY) )
    {
        cbKeySize -= sizeof(AES256KEY);
    }
    else
    {
        cbKeySize = 0;
    }

    BYTE* const pbData = ( cbKeySize != 0 ) ? pKey->pbKey : NULL;
    if ( !CryptExportKey( hKey, NULL, PLAINTEXTKEYBLOB, 0, pbData, &cbKeySize ) )
    {
        if ( GetLastError() == ERROR_MORE_DATA )
        {
            // Add header
            *pcbKeySize = cbKeySize + sizeof(AES256KEY);
            Error( ErrERRCheck( JET_errBufferTooSmall ) );
        }

        Error( ErrOSErrFromWin32Err(GetLastError()) );
    }
    Expected( cbKeySize == 44 );
    // Add header
    *pcbKeySize = cbKeySize + sizeof(AES256KEY);
    // A null key generated success instead of ERROR_MORE_DATA, pretend as if it was ERROR_MORE_DATA.
    if ( pbData == NULL )
    {
        Error( ErrERRCheck( JET_errBufferTooSmall ) );
    }

    pKey->Version = JET_EncryptionAlgorithmAes256;
    pKey->Checksum = Crc32Checksum( pKey->pbKey, cbKeySize );

#ifdef DEBUG
    ULONG blockSize;
    ULONG blocksizeLen = sizeof(blockSize);
    if ( CryptGetKeyParam( hKey, KP_BLOCKLEN, (BYTE*)&blockSize, &blocksizeLen, 0 ) )
    {
        // AES256 uses 16-byte block size
        Expected( blockSize == BlockSizeAes256*8 );
    }

    CallS( ErrOSEncryptionVerifyKey( pbKey, *pcbKeySize ) );
#endif

HandleError:

    CryptDestroyKey( hKey );
    return err;
}

// PERSISTED
struct AES256BLOBTRAILER
{
    BYTE    Version;
    BYTE    InitVector[BlockSizeAes256];
};

ULONG CbOSEncryptAes256SizeNeeded( ULONG cbDataLen )
{
    C_ASSERT( sizeof(AES256BLOBTRAILER) == 1 + BlockSizeAes256 );
    cbDataLen += sizeof(ULONG);                             // CRC32 checksum of data
    return cbDataLen                                        // data
        + ( BlockSizeAes256 - cbDataLen % BlockSizeAes256 ) // AES256 padding
        + sizeof(AES256BLOBTRAILER);                        // version + Init Vector
}

ERR
ErrOSEncryptWithAes256(
    _Inout_updates_bytes_to_(cbDataBufLen, *pcbDataLen)     BYTE *pbData,
    _Inout_                                                 ULONG *pcbDataLen,
    _In_                                                    ULONG cbDataBufLen,
    _In_reads_bytes_(cbKey)                         const   BYTE *pbKey,
    _In_                                                    ULONG cbKey )
{
    ERR err = JET_errSuccess;
    AES256KEY *pKey = (AES256KEY *)pbKey;
    HCRYPTKEY hKey = NULL;
    ULONG checksum;
    AES256BLOBTRAILER trailer;
    ULONG cbNeeded;

    CallR( ErrOSIAESProviderInit() );

    CallR( ErrOSEncryptionVerifyKey( pbKey, cbKey ) );

    cbNeeded = CbOSEncryptAes256SizeNeeded( *pcbDataLen );
    if ( cbNeeded > cbDataBufLen )
    {
        *pcbDataLen = cbNeeded;
        return ErrERRCheck( JET_errBufferTooSmall );
    }

    checksum = Crc32Checksum( pbData, *pcbDataLen );
    // checksum is appended to the plaintext
    *(UnalignedLittleEndian<ULONG> *)(pbData + *pcbDataLen) = checksum;
    *pcbDataLen += sizeof(checksum);

    if ( !CryptImportKey( g_hAESProv, pKey->pbKey, cbKey - sizeof(AES256KEY), NULL, 0, &hKey ) )
    {
        return ErrOSErrFromWin32Err(GetLastError());
    }

    // We are depending on OS default of CBC mode and PKCS5 padding
    if ( !CryptGenRandom( g_hAESProv, sizeof(trailer.InitVector), trailer.InitVector ) ||
         !CryptSetKeyParam( hKey, KP_IV, trailer.InitVector, 0 ) )
    {
        Error( ErrOSErrFromWin32Err(GetLastError()) );
    }

    if ( !CryptEncrypt( hKey, 0, TRUE, 0, pbData, pcbDataLen, cbDataBufLen ) )
    {
        Error( ErrOSErrFromWin32Err(GetLastError()) );
    }

    // Version+InitVector is appended to the ciphertext
    if ( *pcbDataLen + sizeof(trailer) > cbDataBufLen )
    {
        Assert( fFalse ); // Messed up the calculation above?
        *pcbDataLen += sizeof(trailer);
        Error( ErrERRCheck( JET_errBufferTooSmall ) );
    }
    Assert( *pcbDataLen + sizeof(trailer) == cbNeeded );

    trailer.Version = JET_EncryptionAlgorithmAes256;
    memcpy_s( pbData + *pcbDataLen, cbDataBufLen - *pcbDataLen, &trailer, sizeof(trailer) );
    *pcbDataLen += sizeof(trailer);

HandleError:

    Assert( err < JET_errSuccess || *pcbDataLen <= cbDataBufLen );

    if ( hKey )
    {
        CryptDestroyKey( hKey );
    }

    return err;
}

ERR
ErrOSDecryptWithAes256(
    _In_reads_( *pcbDataLen )                           BYTE *pbDataIn,
    _Out_writes_bytes_to_(*pcbDataLen, *pcbDataLen)     BYTE *pbDataOut,
    _Inout_                                             ULONG *pcbDataLen,
    _In_reads_bytes_(cbKey)                     const   BYTE *pbKey,
    _In_                                                ULONG cbKey )
{
    ERR err = JET_errSuccess;
    AES256KEY *pKey = (AES256KEY *)pbKey;
    HCRYPTKEY hKey = NULL;
    ULONG checksum;
    AES256BLOBTRAILER *ptrailer;

    CallR( ErrOSIAESProviderInit() );

    CallR( ErrOSEncryptionVerifyKey( pbKey, cbKey ) );

    if ( *pcbDataLen < BlockSizeAes256 + sizeof(AES256BLOBTRAILER) ||
         *pcbDataLen % BlockSizeAes256 != sizeof(AES256BLOBTRAILER) % BlockSizeAes256 )
    {
        return ErrERRCheck( JET_errInvalidParameter );
    }
    // Version+InitVector is appended to the ciphertext
    ptrailer = (AES256BLOBTRAILER *)(pbDataIn + *pcbDataLen) - 1;
    if ( ptrailer->Version != JET_EncryptionAlgorithmAes256 )
    {
        return ErrERRCheck( JET_errInvalidParameter );
    }
    *pcbDataLen -= sizeof(AES256BLOBTRAILER);

    if ( !CryptImportKey( g_hAESProv, pKey->pbKey, cbKey - sizeof(AES256KEY), NULL, 0, &hKey ) )
    {
        return ErrERRCheck( JET_errInvalidParameter );
    }

    // We are depending on OS default of CBC mode and PKCS5 padding
    if ( !CryptSetKeyParam( hKey, KP_IV, ptrailer->InitVector, 0 ) )
    {
        Error( ErrOSErrFromWin32Err(GetLastError()) );
    }

    // assert that the in/out buffers are not overlapping
    Assert( pbDataOut < pbDataIn || pbDataOut >= ( pbDataIn + *pcbDataLen ) );
    Assert( pbDataIn < pbDataOut || pbDataIn >= ( pbDataOut + *pcbDataLen ) );

    memcpy( pbDataOut, pbDataIn, *pcbDataLen );
    if ( !CryptDecrypt( hKey, NULL, TRUE, 0, pbDataOut, pcbDataLen ) )
    {
        Error( ErrERRCheck( JET_errDecryptionFailed ) );
    }
    if ( *pcbDataLen < sizeof(checksum) )
    {
        Error( ErrERRCheck( JET_errDecryptionFailed ) );
    }
    *pcbDataLen -= sizeof(checksum);
    // checksum is appended to the plaintext
    checksum = *(UnalignedLittleEndian<ULONG> *)(pbDataOut + *pcbDataLen);

    if ( checksum != Crc32Checksum( pbDataOut, *pcbDataLen ) )
    {
        Error( ErrERRCheck( JET_errDecryptionFailed ) );
    }

HandleError:

    if ( hKey )
    {
        CryptDestroyKey( hKey );
    }

    return err;
}

